import os
import json
from pathlib import Path
import torchaudio
import argparse
import torch
from pyarrow.parquet import ParquetDataset

def save_audio(audio_data, output_dir, file_id):
    """保存音频文件并返回路径"""
    os.makedirs(output_dir, exist_ok=True)
    audio_path = os.path.join(output_dir, f"{file_id}.wav")
    
    # 处理包含bytes或path的音频数据
    if isinstance(audio_data, dict):
        if 'bytes' in audio_data:
            # 从bytes加载音频
            import io
            import soundfile as sf
            bytes_io = io.BytesIO(audio_data['bytes'])
            waveform, sampling_rate = sf.read(bytes_io)
            waveform = torch.tensor(waveform)
        elif 'path' in audio_data:
            # 从路径加载音频
            waveform, sampling_rate = torchaudio.load(audio_data['path'])
        else:
            raise ValueError("不支持的音频数据格式")
    else:
        raise ValueError("音频数据必须是字典格式")
    
    if len(waveform.shape) == 1:
        waveform = waveform.unsqueeze(0)
    torchaudio.save(audio_path, waveform, sampling_rate)
    return audio_path

def process_parquet_files(input_dir, output_jsonl, audio_output_dir='output_audio'):
    """
    处理目录中的parquet文件
    :param input_dir: 输入目录或parquet文件路径
    :param output_jsonl: 输出jsonl文件路径
    :param audio_output_dir: 音频输出目录
    """
    results = []
    
    # 获取所有parquet文件
    input_path = Path(input_dir)
    if input_path.is_file() and input_path.suffix == '.parquet':
        parquet_files = [input_path]
    else:
        parquet_files = list(input_path.rglob('*.parquet'))

    for parquet_file in parquet_files:
        try:
            dataset = ParquetDataset(parquet_file)
            table = dataset.read()
            for i, record in enumerate(table.to_pylist()):
                if 'audio' not in record:
                    continue
                    
                try:
                    # 直接传递音频数据，不再转换格式
                    audio_path = save_audio(
                        record['audio'], 
                        audio_output_dir, 
                        f"{parquet_file.stem}_{i}"
                    )
                    # 保留其他数据并添加音频路径
                    item = {k: v for k, v in record.items() if k != 'audio'}
                    item['audio_path'] = os.path.abspath(audio_path)
                    results.append(item)
                
                except Exception as e:
                    print(f"处理 {parquet_file} 第 {i} 条记录失败: {str(e)}")
                    continue
                    
        except Exception as e:
            print(f"读取parquet文件 {parquet_file} 失败: {str(e)}")
            continue

    # 保存为jsonl
    with open(output_jsonl, 'w', encoding='utf-8') as f:
        for record in results:
            f.write(json.dumps(record, ensure_ascii=False) + '\n')

def main():
    parser = argparse.ArgumentParser(description='处理parquet音频数据')
    parser.add_argument('--input', required=True, help='输入目录或parquet文件路径')
    parser.add_argument('--output', default='output.jsonl', help='输出jsonl文件路径')
    parser.add_argument('--audio_dir', default='output_audio', help='音频输出目录')
    args = parser.parse_args()

    process_parquet_files(args.input, args.output, args.audio_dir)
    print(f"处理完成，结果保存在 {os.path.abspath(args.output)}")

if __name__ == '__main__':
    main()